package main

import (
	"context"
	"crypto/tls"
	"errors"
	"fmt"
	"net"
	"shell/common"
	"shell/common/packages"
	"shell/global"
	"shell/utils"
	"sync"
	"sync/atomic"
	"time"

	"github.com/google/uuid"
	"github.com/quic-go/quic-go"
)

// LazyQuicConn 对原始的quic连接进行一次包装
// 实现懒连接：在所有stream都超时时、或quic连接断开时，调用OpenStream会自动重连
// 在没有打开流时，就会延迟quic原始连接的创建
type LazyQuicConn struct {
	conn             *common.QuicConn
	lock             sync.Mutex
	tlsCfg           *tls.Config
	peerAddr         string
	lastStream       atomic.Int64 // 上一次使用OpenStream方法打开流的时间戳
	keepaliveSeconds int64        // 在最后一个流开启后保活的秒数，如果为0，表示不保活，如果为负数表示一直保活
}

// NewLazyQuicConn ks为在最后一个流开启后保活的秒数，如果为0，表示不保活，如果为负数表示一直保活，建议ks为 60 * 5
func NewLazyQuicConn(addr string, tlsCfg *tls.Config, ks int64) (*LazyQuicConn, error) {
	if tlsCfg == nil {
		return nil, global.ErrNil
	}
	_, _, err := net.SplitHostPort(addr)
	if err != nil {
		return nil, err
	}
	result := new(LazyQuicConn)
	result.conn = nil
	result.tlsCfg = tlsCfg
	result.peerAddr = addr
	result.keepaliveSeconds = ks
	result.lastStream.Store(time.Now().Unix())
	return result, nil
}

func (lqc *LazyQuicConn) OpenStream() (*common.QuicStream, error) {
	lqc.lock.Lock()
	defer lqc.lock.Unlock()
	if lqc.conn == nil || !lqc.conn.IsConn.Load() {
		// quic连接没有建立或已经断开
		lqc.conn = nil
		quicConn, err := quic.DialAddr(context.Background(), lqc.peerAddr, lqc.tlsCfg, &quic.Config{
			MaxIdleTimeout:     15 * time.Second,
			MaxIncomingStreams: 1 << 60,
		})
		if err != nil {
			return nil, err
		}
		lqc.conn = common.NewQuicConn(quicConn, bytePool)
		stream, err := lqc.conn.OpenStream()
		if err != nil {
			return stream, err
		}
		lqc.lastStream.Store(time.Now().Unix())
		if lqc.keepaliveSeconds != 0 {
			go func(lazy *LazyQuicConn) {
				ticker := time.NewTicker(time.Second * 5)
				defer ticker.Stop()
				// 向代理端保活
				for {
					select {
					case <-lazy.conn.Ctx().Done():
						return
					case <-ticker.C:
						// 判断是否已经超时
						if lazy.keepaliveSeconds > 0 && lazy.lastStream.Load()+lazy.keepaliveSeconds < time.Now().Unix() {
							return
						}
						s, err := lazy.conn.OpenStream()
						if err != nil {
							common.HandleErr(err)
							return
						}
						req := new(packages.KeepaliveReq)
						req.CmdUUID = uuid.New()
						req.Role = Role
						err = s.WritePackage(req)
						if err != nil {
							common.HandleErr(err)
							s.CloseConn(common.CloseByProcess)
							return
						}
						resp := new(packages.KeepaliveResp)
						err = s.ReadPackage(resp)
						if err != nil {
							common.HandleErr(err)
							s.CloseConn(common.CloseByProcess)
							return
						}
						_ = s.CloseStreamSync(Role, true)
					}
				}
			}(lqc)
		}
		return stream, err
	}
	// quic连接没有断开
	stream, err := lqc.conn.OpenStream()
	if err != nil {
		return stream, err
	}
	lqc.lastStream.Store(time.Now().Unix())
	return stream, nil
}

// 端口转发实例，多线程安全，用于本地转发
type PortRedirInstance struct {
	redirReq     *packages.PortRedirReq
	quicConn     *LazyQuicConn
	pool         *utils.BytesPool
	tcpListener  net.Listener
	isListenning atomic.Bool // 用于
}

// 新建一个端口转发实例
func NewPortRedirInst(req *packages.PortRedirReq, quicConn *LazyQuicConn, pool *utils.BytesPool) *PortRedirInstance {
	if req == nil || quicConn == nil || pool == nil {
		logger.Fatalf(0, "nil error")
	}
	result := new(PortRedirInstance)
	result.quicConn = quicConn
	result.redirReq = req
	result.redirReq.SetRole(Role)
	result.isListenning.Store(false)
	result.pool = pool
	return result
}

// 开始转发，该方法会阻塞
func (pri *PortRedirInstance) StartRedir() error {
	listener, err := net.Listen("tcp", net.JoinHostPort("0.0.0.0", fmt.Sprintf("%d", pri.redirReq.RedirSrcPort)))
	if err != nil {
		return err
	}
	defer func() {
		if pri.isListenning.CompareAndSwap(true, false) {
			_ = listener.Close()
		}
	}()
	pri.tcpListener = listener
	pri.isListenning.Store(true)
	for {
		if !pri.isListenning.Load() {
			return errors.New("listener closed")
		}
		tcpConn, err := listener.Accept()
		if err != nil {
			return err
		}
		go pri.handleRedirConn(tcpConn)
	}
}

// 处理单个转发连接
func (pri *PortRedirInstance) handleRedirConn(conn net.Conn) {
	stream, err := pri.quicConn.OpenStream()
	if err != nil {
		common.HandleErr(err)
		return
	}
	err = stream.WritePackage(pri.redirReq)
	if err != nil {
		common.HandleErr(err)
		return
	}
	resp := new(packages.PortRedirResp)
	err = stream.ReadPackage(resp)
	if err != nil {
		common.HandleErr(err)
	}
	if !resp.IsOk {
		logger.Errorf(0, "error port redir error: %s", string(resp.ErrorInfo))
		_ = stream.CloseStreamSync(Role, true)
		return
	}
	common.NewTimeoutRWCPairWarp(conn, stream, global.DefaultIdleTime, timeoutWatcher).RedirSyncNoKey()
}

// 关闭端口转发实例
func (pri *PortRedirInstance) Stop() {
	pri.isListenning.Store(false)
	_ = pri.tcpListener.Close()
}
